#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
微表情识别系统主执行文件
基于时空特征融合的端到端微表情识别系统
"""

import os
import argparse
import torch
import logging
from datetime import datetime

# 导入自定义模块
from src.data.dataloader import create_dataloaders
from src.models.model_factory import create_model
from src.utils.config import load_config
from src.utils.logger import setup_logger
from src.preprocessing.preprocess import preprocess_dataset
from src.evaluation.metrics import evaluate_model
from src.utils.trainer import Trainer


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='微表情识别系统')
    parser.add_argument('--config', type=str, default='configs/default.yaml',
                        help='配置文件路径')
    parser.add_argument('--mode', type=str, default='train',
                        choices=['train', 'eval', 'demo', 'preprocess'],
                        help='运行模式: 训练/评估/演示/预处理')
    parser.add_argument('--dataset', type=str, default='CASME2',
                        choices=['CASME2', 'SMIC', 'SAMM', 'MEGC2023'],
                        help='数据集名称')
    parser.add_argument('--data_dir', type=str, default='./data',
                        help='数据集根目录')
    parser.add_argument('--output_dir', type=str, default='./output',
                        help='输出目录')
    parser.add_argument('--checkpoint', type=str, default=None,
                        help='模型检查点路径')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                        help='运行设备')
    parser.add_argument('--seed', type=int, default=42,
                        help='随机种子')
    parser.add_argument('--verbose', action='store_true',
                        help='是否显示详细日志')
    return parser.parse_args()


def set_seed(seed):
    """设置随机种子以确保可复现性"""
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 加载配置文件
    config = load_config(args.config)
    
    # 合并命令行参数到配置
    for key, value in vars(args).items():
        if key not in ['config'] and value is not None:
            config[key] = value
    
    # 创建输出目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join(args.output_dir, f"{args.mode}_{args.dataset}_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)
    config['output_dir'] = output_dir
    
    # 设置日志
    log_file = os.path.join(output_dir, f"{args.mode}.log")
    setup_logger(log_file, args.verbose)
    logger = logging.getLogger()
    
    # 设置随机种子
    set_seed(args.seed)
    
    # 根据运行模式执行相应操作
    if args.mode == 'preprocess':
        logger.info("开始数据预处理...")
        preprocess_dataset(config)
    
    elif args.mode == 'train' or args.mode == 'eval':
        # 创建数据加载器
        logger.info(f"正在加载 {args.dataset} 数据集...")
        train_loader, val_loader, test_loader = create_dataloaders(config)
        
        # 创建模型
        logger.info("正在构建模型...")
        model = create_model(config)
        
        # 如果指定了检查点，加载预训练权重
        if args.checkpoint:
            logger.info(f"从检查点加载权重: {args.checkpoint}")
            checkpoint = torch.load(args.checkpoint, map_location=args.device)
            model.load_state_dict(checkpoint['model_state_dict'])
        
        if args.mode == 'train':
            # 创建训练器并开始训练
            logger.info("开始训练过程...")
            trainer = Trainer(model, config, train_loader, val_loader)
            trainer.train()
            
            # 在测试集上评估最佳模型
            logger.info("加载最佳模型并在测试集上评估...")
            best_checkpoint = os.path.join(output_dir, 'checkpoints', 'best_model.pth')
            checkpoint = torch.load(best_checkpoint, map_location=args.device)
            model.load_state_dict(checkpoint['model_state_dict'])
        
        # 评估模型
        logger.info("在测试集上评估模型...")
        metrics = evaluate_model(model, test_loader, config)
        
        # 打印评估结果
        logger.info(f"评估结果: {metrics}")
    
    elif args.mode == 'demo':
        logger.info("演示模式尚未实现")
        # TODO: 实现实时演示功能
    
    logger.info("程序执行完成")


if __name__ == "__main__":
    main() 